# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import getpass
import logging
import subprocess
import sys
import tempfile

from absl import app, flags
import jinja2

from coin.base.datetime_util import iterate_date
from coin.support.feed_tool.feed_cache.app.coin2_fast_feed_runner import (load_feed_writer_config,
                                                                          load_subscribers_setting,
                                                                          get_mea_from_config,
                                                                          get_recipe_from_config,
                                                                          get_groups_from_setting)

FLAGS = flags.FLAGS


class NoJobError(Exception):
  pass


def generate_script_setting():
  machines = [
      "feed-02.ap-northeast-2.aws",
      "feed-05.ap-northeast-1.aws",
      "feed-05.cn-hongkong.aliyun",
      "feed-02.eu-west-1.aws",
      "feed-02.us-east-1.aws",
      "feed-01.ap-northeast-1.aws.huobi",
      "feed-01.ap-southeast-1.aws",
  ]
  config, config_file = load_feed_writer_config(machines)
  subscribers = load_subscribers_setting()
  workers = FLAGS.workers
  settings = {}
  for m in machines:
    if (FLAGS.machine is not None) and (m not in FLAGS.machine):
      continue
    settings[m] = {}
    meas = get_mea_from_config(m, config)
    for mea in meas:
      if (FLAGS.mea is not None) and (mea not in FLAGS.mea) or ('_split_channel' in mea):
        continue
      settings[m][mea] = {}
      recipes = get_recipe_from_config(m, mea, config, subscribers)
      groups = get_groups_from_setting(m, mea, config)
      for recipe in recipes:
        if (FLAGS.recipe is not None) and (recipe not in FLAGS.recipe):
          continue
        settings[m][mea][recipe] = {}
        for group in groups:
          if (FLAGS.group is not None) and (group not in FLAGS.group):
            continue
          settings[m][mea][recipe][group] = {
              "workers": workers,
              "raw": FLAGS.raw,
          }
  return settings


def generate_slurm_script(start_date,
                          end_date,
                          git_commit,
                          summary_mailto,
                          job_submit_period=None,
                          outdir_prefix=None,
                          user=None,
                          grid=None,
                          tag=None,
                          interval=None):
  job_submit_period = 0.01 if job_submit_period is None else job_submit_period
  grid = grid or 'iosg'
  outdir_prefix = \
      outdir_prefix or '/remote/{grid}/home-2/{user}/scratch'.format(grid=grid, user=user)

  out_dir = '%s/feed_stats-${SLURM_JOB_ID}' % outdir_prefix

  start_date = datetime.datetime.strptime(start_date, '%Y%m%d')
  end_date = datetime.datetime.strptime(end_date, '%Y%m%d')
  trading_dates = [td for td in iterate_date(start_date.date(), end_date.date())]
  trading_dates = [td.strftime('%Y%m%d') for td in trading_dates]

  settings = generate_script_setting()

  if tag:
    mail_subject = "[coin-feed-stats] ${SLURM_JOB_ID} summary - %s" % tag
  else:
    mail_subject = "[coin-feed-stats] ${SLURM_JOB_ID} summary"

  template_loader = jinja2.FileSystemLoader(searchpath="./")
  template_env = jinja2.Environment(loader=template_loader)

  template = template_env.get_template(
      'coin/support/feed_tool/feed_stats/app/feed_stats_motion/slurm_script.tmpl')
  return template.render(user=user,
                         out_dir=out_dir,
                         job_submit_period=job_submit_period,
                         trading_dates=trading_dates,
                         settings=settings,
                         git_commit=git_commit,
                         mail_subject=mail_subject,
                         summary_mailto=summary_mailto,
                         interval=interval)


class SBatchError(Exception):
  pass


def run_sbatch(script_filepath, sbatch_args=None, verbose=False):
  sbatch_args = sbatch_args or []
  if isinstance(sbatch_args, str):
    sbatch_args = [token for token in sbatch_args.split(' ') if token]

  sbatch_args = ['--parsable'] + sbatch_args
  if sbatch_args:
    logging.debug('sbatch_args: %s', ' '.join(sbatch_args))

  completed = subprocess.run(['sbatch', *sbatch_args, script_filepath], stdout=subprocess.PIPE)
  if completed.returncode != 0:
    raise SBatchError(completed)

  sbatch_output = completed.stdout.decode('utf-8')
  job_id = int(sbatch_output.strip())
  if verbose:
    sys.stderr.write('Job ID: %d\n' % job_id)
  return job_id


def run_feed_stats_on_slurm(start_date,
                            end_date,
                            git_commit,
                            job_submit_period=None,
                            outdir_prefix=None,
                            user=None,
                            grid=None,
                            venv=None,
                            force_use_venv=True,
                            summary_mailto='{user}@prestolabs.io',
                            tag=None,
                            dry_run=False,
                            sbatch_args=None,
                            notify_user=None,
                            verbose=False,
                            interval=None):
  if force_use_venv and venv is None:
    user = user or getpass.getuser()
    grid = grid or 'iosg'

  if summary_mailto is not None:
    if '{user}' in summary_mailto:
      user = notify_user or 'leon'
      summary_mailto = summary_mailto.format(user=user)
    if verbose:
      sys.stderr.write('Feed stats result mail will be sent to %s.\n' % summary_mailto)

  try:
    script = generate_slurm_script(start_date,
                                   end_date,
                                   git_commit,
                                   summary_mailto,
                                   job_submit_period=job_submit_period,
                                   outdir_prefix=outdir_prefix,
                                   user=getpass.getuser(),
                                   grid=grid,
                                   tag=tag,
                                   interval=interval)
  except NoJobError:
    if verbose:
      sys.stderr.write('No job to run.\n')
    raise

  if dry_run:
    if verbose:
      print(script)
    return

  with tempfile.NamedTemporaryFile(mode='w') as script_file:
    script_file.write(script)
    script_file.flush()

    try:
      job_id = run_sbatch(script_file.name, sbatch_args, verbose=verbose)
      if verbose:
        print(job_id)
    except SBatchError:
      if verbose:
        sys.stderr.write('Error while running sbatch. Abort.\n')
      raise

  return job_id


def main(argv):

  try:
    run_feed_stats_on_slurm(start_date=FLAGS.start_date,
                            end_date=FLAGS.end_date,
                            git_commit=FLAGS.git_commit,
                            outdir_prefix=FLAGS.outdir,
                            grid=FLAGS.grid,
                            summary_mailto=FLAGS.mailto,
                            tag=FLAGS.tag,
                            dry_run=FLAGS.script_only,
                            sbatch_args=FLAGS.sbatch_args,
                            notify_user=FLAGS.notify_user,
                            verbose=True,
                            interval=FLAGS.interval)
  except NoJobError:
    return 0
  except SBatchError:
    return 1


if __name__ == '__main__':
  logging.basicConfig(
      level='DEBUG',
      format='%(levelname)8s %(asctime)s %(name)s %(filename)s:%(lineno)d] %(message)s')
  logging.getLogger('coin.exchange.util.feed_checker').setLevel(logging.WARNING)

  flags.DEFINE_string('start_date', None, 'yyyymmdd')
  flags.DEFINE_string('end_date', None, 'yyyymmdd. Feed at end_date is not checked.')
  flags.DEFINE_string('git_commit', 'HEAD', 'Git commit sha or HEAD.')
  flags.DEFINE_string('outdir',
                      None,
                      'Output directory. Default is /remote/<grid>/home-2/<user>/scratch')
  flags.DEFINE_string('grid', 'iosg', 'Specify slurm grid you will use.')
  flags.DEFINE_string('mailto',
                      '{user}@prestolabs.io',
                      'If specified, it will send a mail after sim finishes.')
  flags.DEFINE_string('tag', None, 'If specified, tag will be in a result mail subject')
  flags.DEFINE_bool('script_only',
                    False,
                    'If true, it will print sbatch script and not run sbatch.')
  flags.DEFINE_string('sbatch_args', '-n 20', 'sbatch arguments')
  flags.DEFINE_string('notify_user', 'leon', 'Email notification receiver.')
  flags.DEFINE_list('machine', None, "Separate by comma")

  flags.DEFINE_list('mea', None, "Separate by comma")

  flags.DEFINE_list('recipe', None, "Separate by comma")

  flags.DEFINE_list('group', None, "Separate by comma")

  flags.DEFINE_list('workers', [0, 1, 2], "Separate by comma")

  flags.DEFINE_bool('raw', False, 'Run feed stats from raw feed.')
  
  flags.DEFINE_string('interval', None, 'interval time for interval feed generate, mmM or hhH')
  app.run(main)
